import random
import re
from itertools import groupby
from typing import List, Set, Tuple

from chernc.compiler.rust.error_messages import RustcErrorMessages, RustcErrorSpan, RustcErrorText
from chernc.core.result import UseManager
from chernc.parser.rust_parser import check_grammar
from chernc.utils.md_utils import extract_code_blocks_with_language


def extract_rust_code_from_md(md_string: str):
    for lang, content in extract_code_blocks_with_language(md_string):
        if lang == "rust" or not check_grammar(content):
            return content
    return None


def extract_toml_code_from_md(md_string: str):
    for lang, content in extract_code_blocks_with_language(md_string):
        if lang == "toml":
            return content
    return None


def error_is_related(msg: RustcErrorMessages, context_end: int, anonymous: bool = True) -> bool:
    if msg.message_type != "diagnostic":
        return False
    if msg.level != "error":
        return False

    def related_span(span: RustcErrorSpan) -> bool:
        while span.expansion:
            span = span.expansion.span
        if span.line_start <= context_end:
            return False
        if anonymous and span.file_name != "<anon>":
            return False
        return True

    return any(related_span(span) for span in msg.spans)


def adjust_msg_line(msg: RustcErrorMessages, context_end: int) -> RustcErrorMessages:
    new_msg = msg.model_copy()
    new_msg.spans = []
    for span in msg.spans:
        while span.expansion:
            span = span.expansion.span
        if span.file_name != "<anon>":
            continue
        # if span.line_start < context_end:
        #     continue
        new_span = span.model_copy()
        new_span.line_start -= context_end
        new_span.line_end -= context_end
        new_msg.spans.append(new_span)
    new_msg.children = [adjust_msg_line(child, context_end) for child in msg.children]

    return new_msg


def format_error_message(msg: RustcErrorMessages, applicability: bool) -> str:
    def highlight_text(text: RustcErrorText) -> str:
        res = ""
        hl_start = text.highlight_start
        hl_len = text.highlight_end - text.highlight_start
        # if hl_len == 0:
        #     return ""
        res += f"`{text.text}`\n"
        res += " " * hl_start + "^" * hl_len + " \n"
        return res

    def format_span(span: RustcErrorSpan) -> str:
        nonlocal applicability
        if span.expansion:
            span = span.expansion.span
        res = ""
        for text in span.text:
            res += highlight_text(text)
        if span.label:
            if res.endswith("\n"):
                res = res[:-1]
            res = res + span.label + "\n"
        if span.suggested_replacement:
            if applicability:
                res += f"suggest({span.suggestion_applicability}): "
            else:
                res += "suggest: "
            if "\n" in span.suggested_replacement:
                res += "\n```\n" + span.suggested_replacement + "```\n"
            else:
                res += "`" + span.suggested_replacement + "`\n"
        if res:
            return f"At line {span.line_start}, col {span.column_start}\n" + res
        else:
            # nothing useful to display
            return ""

    res = f"{msg.level}: {msg.message}\n"
    # if msg.code:
    #     res += f"explanation of error:\n {msg.code['explanation']}\n"
    for span in msg.spans:
        res += format_span(span)
    res += "\n"
    for child in msg.children:
        res += format_error_message(child, applicability)
    return res


def format_error_messages(msgs: List[RustcErrorMessages], applicability=False) -> str:
    return "".join(format_error_message(msg, applicability) for msg in msgs)


def annotate_lines_for_code(code: str, lines: Set[int]) -> str:
    code_lines = code.splitlines()
    for line in lines:
        if line > len(code_lines) or line <= 0:
            continue
        code_lines[line - 1] += f" // line {line}"
    return "\n".join(code_lines)


def filter_error_messages(msgs: List[RustcErrorMessages], anonymous: bool = True) -> List[RustcErrorMessages]:
    msgs = list(filter(lambda m: error_is_related(m, -1, anonymous), msgs))
    if not msgs:
        return []
    else:
        return msgs


def adjust_messages(msgs: List[RustcErrorMessages], context_end: int) -> List[RustcErrorMessages]:
    msgs = list(filter(lambda m: error_is_related(m, context_end), msgs))
    if not msgs:
        return []
    return [adjust_msg_line(msg, context_end) for msg in msgs]


def select_messages(msgs: List[RustcErrorMessages]) -> List[RustcErrorMessages]:
    if not msgs:
        return []
    has_help = [msg for msg in msgs if any(child.level == "help" for child in msg.children)]
    if has_help:
        return has_help[: min(5, len(has_help))]
    errors_of_kind = map(lambda t: (t[0], list(t[1])), groupby(msgs, key=lambda m: m.code and m.code.code))
    _, most_kind = max(errors_of_kind, key=lambda t: len(t[1]))

    return random.sample(most_kind, min(5, len(most_kind)))


def extern_use_rest(code: str) -> Tuple[Set[str], UseManager, str]:
    """
    Split the code into three parts, extern declarations, use statements, and the rest of the code.
    """
    extern_decls = set()
    uses = UseManager()

    def add_extern(m):
        nonlocal extern_decls
        extern_decls.add(m[1])
        return ""

    def add_use(m):
        nonlocal uses
        uses.use(m[1])
        return ""

    code = re.sub(r"extern\s+crate\s+(\w+);", add_extern, code)
    code = re.sub(r"use\s+(\w.+);", add_use, code)
    return extern_decls, uses, code


def into_whole(extern: Set[str], uses: UseManager, code: str) -> str:
    """
    Combine the extern declarations, use statements, and the rest of the code into a whole code.
    """
    res = ""
    for crate in extern:
        res += f"extern crate {crate};\n"
    res += uses.use_declarations() + "\n"
    res += code
    return res
